/*
 * Copyright (c) 2025 lhjhit
 * SPDX-License-Identifier: Apache-2.0
 */

#define DT_DRV_COMPAT ti_drv8889

#include <zephyr/kernel.h>
#include <zephyr/logging/log.h>

#include <zephyr/drivers/gpio.h>
#include <zephyr/drivers/spi.h>
#include <zephyr/drivers/stepper.h>

// #include "step_dir_stepper_common.h"
#include "../../../../zephyr/drivers/stepper/step_dir/step_dir_stepper_common.h"
#include "drv8889.h"

LOG_MODULE_REGISTER(DRV8889_STEPPER_MODULE, LOG_LEVEL_INF);

/* Wake up times of the drv8889 stepper controller family. Only after they have elapsed
 * are controller output signals guaranteed to be valid.
 */
#define DRV8889_WAKE_UP_TIME K_USEC(900)

#define DRV8889_WRITE_CMD 0
#define DRV8889_READ_CMD 1

#define DRV8889_REG_CTRL1 0x03
#define DRV8889_REG_CTRL2 0x04
#define DRV8889_REG_CTRL3 0x05
#define DRV8889_REG_CTRL4 0x06
#define DRV8889_REG_CTRL5 0x07
#define DRV8889_REG_CTRL6 0x08
#define DRV8889_REG_CTRL7 0x09
#define DRV8889_REG_CTRL8 0x0A

#define CHIPID 0x02

typedef union {
	uint16_t two_bytes;
	struct {
		uint8_t low_byte;
		uint8_t high_byte;
	} one_byte;
} u16_u8_data;


/**
 * @brief DRV8889 stepper driver configuration data.
 *
 * This structure contains all of the devicetree specifications for the pins
 * needed by a given DRV8889 stepper driver.
 */
struct drv8889_config {
	struct step_dir_stepper_common_config common;
	struct gpio_dt_spec nsleep_pin;
	struct gpio_dt_spec drvoff_pin;
	struct gpio_dt_spec nfault_pin;
	struct spi_dt_spec spi;
};

/**
 * @brief DRV8889 stepper driver data.
 */
struct drv8889_data {
	const struct step_dir_stepper_common_data common;
	const struct device *dev;
	enum stepper_micro_step_resolution ustep_res;
	struct gpio_callback fault_cb_data;
};

static int drv8889_reg_read(const struct device *dev, uint8_t addr, uint16_t *value)
{
	const struct drv8889_config *config = dev->config;
	uint16_t write_data = (DRV8889_READ_CMD << 14) | (addr << 9);
	const struct spi_buf tx_buf = { .buf = &write_data, .len = sizeof(write_data) };
	struct spi_buf rx_buf = { .buf = value, .len = sizeof(uint16_t) };
	const struct spi_buf_set tx = { .buffers = &tx_buf, .count = 1 };
	const struct spi_buf_set rx = { .buffers = &rx_buf, .count = 1 };

	return spi_transceive_dt(&config->spi, &tx, &rx);
}

static int drv8889_get_chipid(const struct device *dev, uint8_t *value)
{
	u16_u8_data read_data;
	int ret;

	ret = drv8889_reg_read(dev, DRV8889_REG_CTRL8, &read_data.two_bytes);
	if (ret < 0) {
		LOG_ERR("%s: Failed to read data (error: %d)", dev->name, ret);
		return ret;
	}
	
	*value = read_data.one_byte.low_byte;

	return 0;
}

static int drv8889_enable(const struct device *dev)
{
	const struct drv8889_config *config = dev->config;
	struct drv8889_data *data = dev->data;
	bool has_drvoff_pin = config->drvoff_pin.port != NULL;
	bool has_nsleep_pin = config->nsleep_pin.port != NULL;
	bool has_nfault_pin = config->nfault_pin.port != NULL;
	k_timeout_t enable_timeout;
	int ret;
	
	if (has_drvoff_pin) {
		ret = gpio_pin_set_dt(&config->drvoff_pin, false);
		if (ret != 0) {
			LOG_ERR("%s: Failed to set drvoff_pin (error: %d)", dev->name, ret);
			return ret;
		}
	}

	if (has_nsleep_pin) {
		ret = gpio_pin_set_dt(&config->nsleep_pin, false);
		if (ret != 0) {
			LOG_ERR("%s: Failed to set nsleep_pin (error: %d)", dev->name, ret);
			return ret;
		}
	}

	if (has_nsleep_pin | has_drvoff_pin) {
		enable_timeout = DRV8889_WAKE_UP_TIME;
	}

	if (has_nfault_pin) {
		/* Wait after enable/wakeup until the fault pin is guaranteed to be in the
		 * proper state.
		 */
		k_sleep(enable_timeout);
		ret = gpio_add_callback_dt(&config->nfault_pin, &data->fault_cb_data);
		if (ret != 0) {
			LOG_ERR("%s: Failed to add fault callback (error: %d)", dev->name, ret);
			return ret;
		}
	}

	return 0;
}

static int drv8889_disable(const struct device *dev)
{
	const struct drv8889_config *config = dev->config;
	struct drv8889_data *data = dev->data;
	bool has_drvoff_pin = config->drvoff_pin.port != NULL;
	bool has_nsleep_pin = config->nsleep_pin.port != NULL;
	bool has_nfault_pin = config->nfault_pin.port != NULL;
	int ret;

	if (has_drvoff_pin) {
		ret = gpio_pin_set_dt(&config->drvoff_pin, true);
		if (ret < 0) {
			LOG_ERR("%s: Failed to set drvoff_pin (error: %d)", dev->name, ret);
			return ret;
		}
	}

	if (has_nsleep_pin) {
		ret = gpio_pin_set_dt(&config->nsleep_pin, true);
		if (ret < 0) {
			LOG_ERR("%s: Failed to set nsleep_pin (error: %d)", dev->name, ret);
			return ret;
		}
	}
	
	if (has_nfault_pin) {
		ret = gpio_remove_callback_dt(&config->nfault_pin, &data->fault_cb_data);
		if (ret < 0) {
			LOG_ERR("%s: Failed to remove fault callback (error: %d)", dev->name, ret);
			return ret;
		}
	}

	return 0;
}

static int drv8889_reg_write(const struct device *dev, uint8_t addr, uint8_t value)
{
	const struct drv8889_config *config = dev->config;
	uint16_t write_data = (DRV8889_WRITE_CMD << 14) | (addr << 9) | value;
	uint16_t read_data = 0;
	const struct spi_buf tx_buf = { .buf = &write_data, .len = sizeof(write_data) };
	const struct spi_buf rx_buf = { .buf = &read_data, .len = sizeof(read_data) };
	const struct spi_buf_set tx = { .buffers = &tx_buf, .count = 1 };
	const struct spi_buf_set rx = { .buffers = &rx_buf, .count = 1 };

	return spi_transceive_dt(&config->spi, &tx, &rx);
}

/*******************************************************************************
//  TRQ_DAC = 2.2*Ifull/Vref
//  If Ifull is 1A and Vref is 2.2V, then TRQ_DAC is 1(100%)
//	TRQ_DAC[7:4]
//		0000b				100%
//		0001b				93.75%
//		0010b				87.5%
//		0011b				81.25%
//		0100b				75%
//		...
//		1110b				12.5%
//		1111b				6.25%
********************************************************************************/

int drv8889_set_trq_dac(const struct device *dev, uint8_t value)
{
	uint8_t reg_value;
	int ret;

	reg_value = value << 4;

	ret = drv8889_reg_write(dev, DRV8889_REG_CTRL1, reg_value);
	if (ret < 0) {
		LOG_ERR("Failed to set trq_dac: %d", ret);
		return ret;
	}

	return 0;
}

static int drv8889_set_micro_step_res(const struct device *dev,
				      enum stepper_micro_step_resolution micro_step_res)
{
	struct drv8889_data *data = dev->data;
	uint8_t reg_value;
	int ret;

	/*DO NOT support non-circular 1/2 step*/
	switch (micro_step_res) {
	case STEPPER_MICRO_STEP_1:
		reg_value = 0;
		break;
	case STEPPER_MICRO_STEP_2:
		reg_value = 3;
		break;
	case STEPPER_MICRO_STEP_4:
		reg_value = 4;
		break;
	case STEPPER_MICRO_STEP_8:
		reg_value = 5;
		break;
	case STEPPER_MICRO_STEP_16:
		reg_value = 6;
		break;
	case STEPPER_MICRO_STEP_32:
		reg_value = 7;
		break;
	case STEPPER_MICRO_STEP_64:
		reg_value = 8;
		break;
	case STEPPER_MICRO_STEP_128:
		reg_value = 9;
		break;
	case STEPPER_MICRO_STEP_256:
		reg_value = 10;
		break;
	default:
		return -ENOTSUP;
	};

	ret = drv8889_reg_write(dev, DRV8889_REG_CTRL3, reg_value);

	if (ret < 0) {
		LOG_ERR("Failed to set microstep: %d", ret);
		return ret;
	}

	data->ustep_res = micro_step_res;

	return 0;
}

static int drv8889_get_micro_step_res(const struct device *dev,
				      enum stepper_micro_step_resolution *micro_step_res)
{
	struct drv8889_data *data = dev->data;
	*micro_step_res = data->ustep_res;
	return 0;
}

static void fault_event(const struct device *dev, struct gpio_callback *cb, uint32_t pins)
{
	struct drv8889_data *data = CONTAINER_OF(cb, struct drv8889_data, fault_cb_data);

	stepper_trigger_callback(data->dev, STEPPER_EVENT_FAULT_DETECTED);
}

static int drv8889_init(const struct device *dev)
{
	const struct drv8889_config *const config = dev->config;
	struct drv8889_data *const data = dev->data;
	bool has_drvoff_pin = config->drvoff_pin.port != NULL;
	bool has_nsleep_pin = config->nsleep_pin.port != NULL;
	bool has_nfault_pin = config->nfault_pin.port != NULL;
	uint8_t value;
	int ret;

	if (!spi_is_ready_dt(&config->spi)) {
		LOG_ERR("SPI bus is not ready");
		return -ENODEV;
	}

	if (has_nsleep_pin) {
		ret = gpio_pin_configure_dt(&config->nsleep_pin, GPIO_OUTPUT_ACTIVE);
		if (ret < 0) {
			LOG_ERR("%s: Failed to configure nsleep_pin (error: %d)", dev->name, ret);
			return ret;
		}
	}

	if (has_drvoff_pin) {
		ret = gpio_pin_configure_dt(&config->drvoff_pin, GPIO_OUTPUT_INACTIVE);
		if (ret < 0) {
			LOG_ERR("%s: Failed to configure drvoff_pin (error: %d)", dev->name, ret);
			return ret;
		}
	}

	/* Check chip ID, it will return if the chip is not DRV8889 */
	ret = drv8889_get_chipid(dev, &value);
	if(ret < 0) {
		LOG_ERR("Failed to get chipid (error: %d)", ret);
		return ret;
	} else if(value != CHIPID) {
		LOG_ERR("Chip ID is not correct: %x", value);
		return -1;//TODO(lhj): need to choose a error code
	}
	LOG_INF("Chip ID is %x", value);

	ret = drv8889_set_micro_step_res(dev, data->ustep_res);
	if (ret < 0) {
		LOG_ERR("Failed to initialize microstep (error: %d)", ret);
		return ret;
	}

	ret = step_dir_stepper_common_init(dev);
	if (ret < 0) {
		LOG_ERR("Failed to initialize common step direction stepper (error: %d)", ret);
		return ret;
	}

	if (has_nfault_pin) {
		ret = gpio_pin_configure_dt(&config->nfault_pin, GPIO_INPUT);
		if (ret < 0) {
			LOG_ERR("%s: Failed to configure fault_pin (error: %d)", dev->name, ret);
			return ret;
		}

		ret = gpio_pin_interrupt_configure_dt(&config->nfault_pin,
						      GPIO_INT_EDGE_TO_ACTIVE);
		if (ret < 0) {
			LOG_ERR("Error %d: failed to configure interrupt on %s pin %d", ret,
				config->nfault_pin.port->name, config->nfault_pin.pin);
			return ret;
		}

		gpio_init_callback(&data->fault_cb_data, fault_event, BIT(config->nfault_pin.pin));
	}

	return 0;
}

static DEVICE_API(stepper, drv8889_stepper_api) = {
	.enable = drv8889_enable,
	.disable = drv8889_disable,
	.move_by = step_dir_stepper_common_move_by,
	.move_to = step_dir_stepper_common_move_to,
	.is_moving = step_dir_stepper_common_is_moving,
	.set_reference_position = step_dir_stepper_common_set_reference_position,
	.get_actual_position = step_dir_stepper_common_get_actual_position,
	.set_microstep_interval = step_dir_stepper_common_set_microstep_interval,
	.run = step_dir_stepper_common_run,
	.stop = step_dir_stepper_common_stop,
	.set_micro_step_res = drv8889_set_micro_step_res,
	.get_micro_step_res = drv8889_get_micro_step_res,
	.set_event_callback = step_dir_stepper_common_set_event_callback,
};

#define DRV8889_DEVICE(inst)                                                                       \
                                                                                                   \
	static const struct drv8889_config drv8889_config_##inst = {                               \
		.common = STEP_DIR_STEPPER_DT_INST_COMMON_CONFIG_INIT(inst),                       \
		.nsleep_pin = GPIO_DT_SPEC_INST_GET_OR(inst, nsleep_gpios, {0}),                     \
		.drvoff_pin = GPIO_DT_SPEC_INST_GET_OR(inst, drvoff_gpios, {0}),                           \
		.nfault_pin = GPIO_DT_SPEC_INST_GET_OR(inst, nfault_gpios, {0}),                     \
		.spi = SPI_DT_SPEC_INST_GET(                                                       \
			inst,                                                                      \
			(SPI_OP_MODE_MASTER | SPI_WORD_SET(16) | SPI_TRANSFER_MSB | SPI_MODE_CPHA)),  \
	};                                                                                         \
                                                                                                   \
	static struct drv8889_data drv8889_data_##inst = {                                         \
		.common = STEP_DIR_STEPPER_DT_INST_COMMON_DATA_INIT(inst),                         \
		.ustep_res = DT_INST_PROP(inst, micro_step_res),                                   \
		.dev = DEVICE_DT_INST_GET(inst),                                                   \
	};                                                                                         \
                                                                                                   \
	DEVICE_DT_INST_DEFINE(inst, &drv8889_init, NULL, &drv8889_data_##inst,                     \
			      &drv8889_config_##inst, POST_KERNEL, CONFIG_STEPPER_INIT_PRIORITY,   \
			      &drv8889_stepper_api);

DT_INST_FOREACH_STATUS_OKAY(DRV8889_DEVICE)